import backtrader as bt

from strategy.utils.BaoStockPandasData import BaoStockPandasData
from utils.DataSource import get_stock_data2

class MomentumStrategy(bt.Strategy):
    params = (
        ('period', 10),  # 动量计算周期
        ('buy_threshold', 0),  # 买入阈值，正数表示仅在动量大于该值时买入
        ('sell_threshold', 0),  # 卖出阈值，负数表示仅在动量小于该值时卖出
    )

    def __init__(self):
        self.data_close = self.datas[0].close
        self.prev_close = None
        self.order = None

    def next(self):
        if self.order:
            return  # 如果有订单在执行，则不执行新的买卖操作

        # 计算动量
        momentum = self.data_close[0] - (self.prev_close if self.prev_close is not None else self.data_close[0])

        # 买入条件：动量大于买入阈值且无持仓
        if momentum > self.params.buy_threshold and not self.position:
            amount_to_invest = (0.95 * self.broker.cash) / self.data_close[0]
            self.buy(size=amount_to_invest)
            self.log(f"BUY EXECUTED --- Price: {self.data_close[0]}, Momentum: {momentum}")

        # 卖出条件：动量小于卖出阈值且有持仓
        elif momentum < self.params.sell_threshold and self.position:
            self.sell()
            self.log(f"SELL EXECUTED --- Price: {self.data_close[0]}, Momentum: {momentum}")

        self.prev_close = self.data_close[0]

    def log(self, txt, dt=None):
        ''' Logging function for this strategy'''
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    # notify_order等其他方法可以参考之前的示例，这里省略以保持简洁

# 回测设置示例
if __name__ == '__main__':
    cerebro = bt.Cerebro()

    # 假设已经通过某种方式（如Yahoo Finance）获取了数据
    df = get_stock_data2("sh.600000", "2020-01-01", "2023-12-31")
    data = BaoStockPandasData(dataname=df)
    cerebro.adddata(data)

    cerebro.addstrategy(MomentumStrategy, period=10, buy_threshold=1, sell_threshold=-1)

    cerebro.broker.setcash(100000.0)
    cerebro.broker.set_coc(True)  # 现金交易模式

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    cerebro.run()
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
    
    #动量策略基于资产价格的变动趋势，当价格呈现强劲上升或下降趋势时，跟随趋势进行交易。以下是在Backtrader框架中实现一个简单动量策略的示例代码。该策略利用价格动量（通常是收盘价之差）来决定买入或卖出。
    #此策略中，period参数用于定义计算动量的周期长度，buy_threshold和sell_threshold分别设置了买入和卖出的动量阈值。当价格动量超过买入阈值时，策略执行买入操作；当动量低于卖出阈值时，若已有持仓，则执行卖出操作。请根据实际市场情况调整参数。